import numpy as np
from astropy.io import fits
import argparse
import os

def read_binary_file(filename, dtype=np.uint16):
    """从二进制文件中读取数据并转换为 NumPy 数组"""
    with open(filename, 'rb') as file:
        data = np.fromfile(file, dtype=dtype)
    return data

def split_data_into_columns(data, num_columns=16):
    """将数据按列分割成多个子数组"""
    if len(data) % num_columns != 0:
        raise ValueError("数据长度必须是列数的整数倍")
    columns = []
    for i in range(num_columns):
        column = data[i::num_columns]
        columns.append(column.reshape(1050, 1100))  # 重塑为1050×1100图像
    return columns

def average_adjacent_columns(columns):
    """将相邻的两列图像进行平均处理（使用完整的16列）"""
    if len(columns) != 16:
        raise ValueError("必须提供16列图像数据进行处理")
    
    averaged_columns = []
    
    # 处理第一行的8列(索引0-7)
    for i in range(0, 8, 2):
        avg = (columns[i] + columns[i+1]) / 2
        averaged_columns.append(avg)
    
    # 处理第二行的8列(索引8-15)
    for i in range(8, 16, 2):
        avg = (columns[i] + columns[i+1]) / 2
        averaged_columns.append(avg)
    
    return averaged_columns

def concatenate_columns(columns, num_rows, num_columns_per_row):
    """将列数据拼接为二维数组"""
    if len(columns) != num_rows * num_columns_per_row:
        raise ValueError("列的数量必须是每行列数的整数倍")
    rows = []
    for i in range(num_rows):
        start = i * num_columns_per_row
        end = start + num_columns_per_row
        row = np.hstack(columns[start:end])  # 水平拼接每行的列
        rows.append(row)
    combined_image = np.vstack(rows)  # 垂直拼接各行
    return combined_image

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="处理二进制文件并生成筛选后的原始图像和平均后的FITS图像")
    parser.add_argument("bin_file", help="输入的二进制文件路径")
    parser.add_argument("output_dir", help="输出目录")
    parser.add_argument("--selected-output", default="selected_original.fits", 
                       help="筛选后的原始FITS文件名（默认：selected_original.fits）")
    parser.add_argument("--averaged-output", default="averaged_image.fits",
                       help="平均后的FITS文件名（默认：averaged_image.fits）")
    
    args = parser.parse_args()
    
    # 确保输出目录存在
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)
    
    # 1. 读取数据
    data = read_binary_file(args.bin_file)
    
    # 2. 分割数据为16列
    all_columns = split_data_into_columns(data)
    
    # 3. 生成筛选后的原始图像（只保留第1、3、5、7 9 11 13 15列）
    selected_indices = [2, 3, 6, 7, 8, 9, 12, 13]  # 对应第1、3、5、7,9 11 13 15列（Python从0开始计数）
    # selected_indices = [0, 4, 8, 12] 
    selected_columns = [all_columns[i] for i in selected_indices]
    selected_image = concatenate_columns(selected_columns, num_rows=2, num_columns_per_row=4)
    selected_path = os.path.join(args.output_dir, args.selected_output)
    fits.PrimaryHDU(data=selected_image).writeto(selected_path, overwrite=True)
    print(f"筛选后的原始FITS文件已生成：{selected_path}")
    
    # 4. 生成平均后的2行4列图像（使用完整的16列）
    averaged_columns = average_adjacent_columns(all_columns)  # 使用全部16列进行平均
    averaged_image = concatenate_columns(averaged_columns, num_rows=2, num_columns_per_row=4)
    averaged_path = os.path.join(args.output_dir, args.averaged_output)
    fits.PrimaryHDU(data=averaged_image).writeto(averaged_path, overwrite=True)
    print(f"平均后的2行4列FITS文件已生成：{averaged_path}")
